package top.jiangqiang.qianyi.common.storage.impl;

import jakarta.annotation.PreDestroy;
import lombok.Getter;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.dromara.hutool.core.io.IoUtil;
import org.dromara.hutool.core.io.file.FileNameUtil;
import org.dromara.hutool.core.io.file.FileUtil;
import org.dromara.hutool.core.text.StrUtil;
import org.dromara.hutool.core.thread.ExecutorBuilder;
import org.dromara.hutool.core.thread.NamedThreadFactory;
import org.springframework.web.util.UriUtils;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.core.ResponseInputStream;
import software.amazon.awssdk.core.async.AsyncRequestBody;
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
import software.amazon.awssdk.http.SdkHttpResponse;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3AsyncClient;
import software.amazon.awssdk.services.s3.S3Configuration;
import software.amazon.awssdk.services.s3.model.*;
import software.amazon.awssdk.transfer.s3.S3TransferManager;
import software.amazon.awssdk.transfer.s3.model.UploadRequest;
import software.amazon.awssdk.transfer.s3.progress.TransferListener;
import top.jiangqiang.qianyi.common.constants.http.Error500Status;
import top.jiangqiang.qianyi.common.exception.Error500Exception;
import top.jiangqiang.qianyi.common.id.IDGenerator;
import top.jiangqiang.qianyi.common.storage.AbstractStorageService;
import top.jiangqiang.qianyi.common.storage.ProgressContext;
import top.jiangqiang.qianyi.common.storage.config.S3StorageConfig;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * @Date 2024/7/5 下午5:20
 * @Author Jiangqiang
 */
@Getter
@Slf4j
public class S3StorageService extends AbstractStorageService<S3StorageConfig> {
    /**
     * 临时上传目录
     */
    private final File tmpDir;
    //线程池，用于异步上传
    private final ExecutorService executorService = ExecutorBuilder.of()
            .setCorePoolSize(3)
            .setMaxPoolSize(100)
            .setThreadFactory(new NamedThreadFactory("s3-upload-thread-pool", true))
            .build();
    //异步上传客户端
    protected S3AsyncClient asyncClient;
    //支持监控传输进度
    protected S3TransferManager transferManager;

    public S3StorageService(IDGenerator idGenerator, S3StorageConfig storageConfig, String cachePath) {
        super(idGenerator, storageConfig, cachePath);
        tmpDir = FileUtil.file(cachePath, "s3");
        if (!tmpDir.exists()) {
            FileUtil.mkdir(tmpDir);
        }
        this.asyncClient = initAsyncClient();
        this.transferManager = S3TransferManager.builder().s3Client(asyncClient).build();
    }

    @Override
    public String generateFileKey(String originalFilename) {
        S3StorageConfig config = getStorageConfig();
        for (; ; ) {
            String fileKey = idGenerator.nextUUID() + "." + FileNameUtil.extName(originalFilename);
            String dirName = config.getDirName();
            if (StrUtil.isNotBlank(dirName)) {
                if (dirName.endsWith("/")) {
                    fileKey = dirName + fileKey;
                } else {
                    fileKey = dirName + "/" + fileKey;
                }
            }
            if (fileKey.startsWith("/")) {
                //不能以斜杠开头，否则会报错
                fileKey = fileKey.substring(1);
            }
            HeadObjectRequest headObjectRequest = HeadObjectRequest.builder()
                    .bucket(config.getBucket())
                    .key(fileKey)
                    .build();
            //判断文件是否存在
            CompletableFuture<Boolean> exceptionally = asyncClient.headObject(headObjectRequest)
                    .thenApply(headObjectResponse -> true)
                    .exceptionally(throwable -> {
                        if (throwable.getCause() instanceof NoSuchKeyException) {
                            return false;
                        } else {
                            log.debug(throwable.getMessage());
                            throw new Error500Exception(Error500Status.INTERNAL_SERVER_ERROR, throwable);
                        }
                    });
            if (!exceptionally.join()) {
                return fileKey;
            }
        }
    }

    @Override
    @SneakyThrows
    public String upload(String originalFilename, InputStream inputStream, Consumer<ProgressContext> progressCallback) {
        String fileKey = generateFileKey(originalFilename);
        S3StorageConfig config = getStorageConfig();
        // 创建 PutObjectRequest 请求
        PutObjectRequest putObjectRequest = PutObjectRequest.builder()
                .bucket(config.getBucket())
                .key(fileKey)
                .build();
        long size = inputStream.available();
        AsyncRequestBody requestBody = AsyncRequestBody.fromInputStream(inputStream, size, executorService);
        UploadRequest uploadRequest = UploadRequest.builder()
                .requestBody(requestBody)
                .putObjectRequest(putObjectRequest)
                .addTransferListener(new TransferListener() {
                    @Override
                    public void transferFailed(Context.TransferFailed context) {
                        //程序出错时会执行这里
                        log.error("文件上传失败", context.exception());
                    }

                    @Override
                    public void transferComplete(Context.TransferComplete context) {
                        context.progressSnapshot().ratioTransferred().ifPresent(progress -> {
                            if (progressCallback != null) {
                                ProgressContext progressContext = new ProgressContext();
                                progressContext.setProgress(progress);
                                progressContext.setComplete(true);
                                progressContext.setFileKey(fileKey);
                                progressCallback.accept(progressContext);
                            }
                        });
                    }

                    @Override
                    public void bytesTransferred(Context.BytesTransferred context) {
                        context.progressSnapshot().ratioTransferred().ifPresent(progress -> {
                            if (progressCallback != null) {
                                ProgressContext progressContext = new ProgressContext();
                                progressContext.setProgress(progress);
                                progressContext.setComplete(false);
                                progressContext.setFileKey(fileKey);
                                progressCallback.accept(progressContext);
                            }
                        });
                    }

                    @Override
                    public void transferInitiated(Context.TransferInitiated context) {
                        //首先会调用这里
                        context.progressSnapshot().ratioTransferred().ifPresent(progress -> {
                            if (progressCallback != null) {
                                ProgressContext progressContext = new ProgressContext();
                                progressContext.setProgress(progress);
                                progressContext.setComplete(false);
                                progressContext.setFileKey(fileKey);
                                progressCallback.accept(progressContext);
                            }
                        });
                    }
                }).build();
        transferManager.upload(uploadRequest)
                .completionFuture()
                .join();
        return fileKey;
    }

    @Override
    public void download(String fileKey, OutputStream outputStream) {
        S3StorageConfig config = getStorageConfig();
        GetObjectRequest getObjectRequest = GetObjectRequest.builder()
                .bucket(config.getBucket())
                .key(fileKey)
                .build();
        CompletableFuture<ResponseInputStream<GetObjectResponse>> responseFuture = asyncClient.getObject(getObjectRequest, AsyncResponseTransformer.toBlockingInputStream());
        try (ResponseInputStream<GetObjectResponse> responseStream = responseFuture.join()) {
            responseStream.transferTo(outputStream); // BLOCKS the calling thread
        } catch (IOException e) {
            throw new Error500Exception(Error500Status.INTERNAL_UPLOAD_FAIL);
        }
    }

    @Override
    public byte[] download(String fileKey) {
        S3StorageConfig config = getStorageConfig();
        GetObjectRequest getObjectRequest = GetObjectRequest.builder()
                .bucket(config.getBucket())
                .key(fileKey)
                .build();
        return asyncClient.getObject(getObjectRequest, AsyncResponseTransformer.toBytes()).join().asByteArray();
    }

    @Override
    public boolean delete(String fileKey) {
        DeleteObjectRequest deleteObjectRequest = DeleteObjectRequest.builder()
                .bucket(getStorageConfig().getBucket())
                .key(fileKey)
                .build();
        SdkHttpResponse sdkHttpResponse = asyncClient.deleteObject(deleteObjectRequest).join().sdkHttpResponse();
        return sdkHttpResponse.isSuccessful();
    }

    @Override
    public boolean test() {
        try {
            this.asyncClient.listBuckets();
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    @Override
    public String getUrl(String fileKey) {
        String objectURL;
        if (StrUtil.isBlank(storageConfig.getDomain())) {
            String host;
            if (storageConfig.getEnablePathStyleAccess()) {
                host = storageConfig.getEndpoint() + "/" + storageConfig.getBucket();
            } else {
                host = storageConfig.getBucket() + "." + storageConfig.getEndpoint();
            }
            objectURL = storageConfig.getAccessProtocol().getProtocolName() + "://" + host + "/" + fileKey;
        } else {
            objectURL = storageConfig.getAccessProtocol().getProtocolName() + "://" + storageConfig.getDomain() + "/" + fileKey;
        }
        return UriUtils.encodePath(objectURL, StandardCharsets.UTF_8);
    }

    @PreDestroy
    @Override
    public void cleanup() {
        // 等待线程池终止
        try {
            executorService.shutdown();
            if (!executorService.isShutdown()) {
                if (!executorService.awaitTermination(3, TimeUnit.SECONDS)) {
                    // 如果超时仍未终止，可以调用 shutdownNow() 方法进行强制关闭
                    executorService.shutdownNow();
                }
            }
        } catch (Exception e) {
            log.debug("关闭线程池出现异常：{}", e.getMessage());
        }
        IoUtil.closeQuietly(this.transferManager, this.asyncClient);
    }

    private S3AsyncClient initAsyncClient() {
        S3StorageConfig config = getStorageConfig();
        return S3AsyncClient.builder()
                .region(Region.of(config.getRegion()))
                .endpointOverride(URI.create(config.getEndpointProtocol() + "://" + config.getEndpoint()))
                .credentialsProvider(() -> AwsBasicCredentials.create(config.getAccessKey(), config.getAccessSecret()))
                .serviceConfiguration(S3Configuration.builder()
                        .chunkedEncodingEnabled(false)
                        .pathStyleAccessEnabled(config.getEnablePathStyleAccess())
                        .build())
                .build();
    }
}
